import numpy as np

def parse_image_file_data(image_file_data: bytes) -> np.array:
    total_images = int.from_bytes(image_file_data[4: 8], byteorder='big')
    n = int.from_bytes(image_file_data[8: 12], byteorder='big')
    m = int.from_bytes(image_file_data[12: 16], byteorder='big')
    images = [0 for _ in range(total_images)]
    for i in range(total_images):
        offset = 16 + n * m * i
        images[i] = 255 - np.frombuffer(image_file_data[offset: offset + n * m], dtype=np.uint8)
        images[i] = images[i].reshape((n, m))
    return images

def parse_label_file_data(label_file_data: bytes):
    total_labels = int.from_bytes(label_file_data[4: 8], byteorder='big')
    labels = list(map(lambda x: x, label_file_data[8: 8 + total_labels]))
    return labels

def read_images_and_labels(image_file_path: str, label_file_path: str):
    image_file = open(image_file_path, "rb")
    image_file_data = image_file.read()
    images = parse_image_file_data(image_file_data)
    label_file = open(label_file_path, "rb")
    label_file_data = label_file.read()
    labels = parse_label_file_data(label_file_data)
    return (images, labels)

class DataLoader:
    def __init__(self, folder_path: str):
        images, labels = read_images_and_labels(folder_path + "/images", folder_path + "/labels")
        self.total = len(images)
        height, width = images[0].shape
        self.images = np.array(images).reshape((self.total, height, width, 1))
        self.labels = np.zeros((self.total, 10))
        self.labels[np.arange(self.total), labels] = 1
    
    def load(self, offset: int, size: int) -> (np.array, np.array):
        ret_images = self.images[offset: offset + size]
        ret_labels = self.labels[offset: offset + size]
        return (ret_images, ret_labels)

    def load_all(self) -> (np.array, np.array):
        return (self.images, self.labels)